{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_id             57080f973d10a288385cf4ca\n",
       "close                            2781.02\n",
       "date                          2016-02-04\n",
       "high                              2793.3\n",
       "low                              2751.31\n",
       "ma10                             2769.24\n",
       "ma20                             2877.77\n",
       "ma5                              2739.26\n",
       "open                             2751.43\n",
       "p_change                            1.52\n",
       "price_change                       41.77\n",
       "v_ma10                       1.73044e+06\n",
       "v_ma20                       1.94013e+06\n",
       "v_ma5                        1.63983e+06\n",
       "volume                       1.70382e+06\n",
       "Name: 1, dtype: object"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pymongo import MongoClient\n",
    "import pandas as pd\n",
    "import datetime\n",
    "\n",
    "# Open Database and find history data collection\n",
    "client = MongoClient()\n",
    "db = client.test_database\n",
    "shdaily = db.indexdata\n",
    "\n",
    "# KDJ calculation formula\n",
    "def KDJCalculation(K1, D1, high, low, close):\n",
    "    # input last K1, D1, max value, min value and current close value\n",
    "    #设定KDJ基期值\n",
    "    #count = 9\n",
    "    #设定k、d平滑因子a、b，不过目前已经约定俗成，固定为1/3\n",
    "    a = 1.0/3\n",
    "    b = 1.0/3\n",
    "    # 取得过去count天的最低价格\n",
    "    low_price = low #low.min() #min(list1)\n",
    "    # 取得过去count天的最高价格\n",
    "    high_price = high #high.max() #max(list1)\n",
    "    # 取得当日收盘价格\n",
    "    current_close = close\n",
    "    if high_price!=low_price:\n",
    "        #计算未成熟随机值RSV(n)＝（Ct－Ln）/（Hn-Ln）×100\n",
    "        RSV = (current_close-low_price)/(high_price-low_price)*100\n",
    "    else:\n",
    "        RSV = 50\n",
    "    #当日K值=(1-a)×前一日K值+a×当日RSV\n",
    "    K2=(1-a)*K1+a*RSV\n",
    "    #当日D值=(1-a)×前一日D值+a×当日K值\n",
    "    D2=(1-b)*D1+b*K2\n",
    "    #计算J值\n",
    "    J2 = 3*K2-2*D2\n",
    "    #log.info(\"Daily K1: %s, D1: %s, K2: %s, D2: %s, J2: %s\" % (K1,D1,K2,D2,J2))\n",
    "    return K1,D1,K2,D2,J2\n",
    "\n",
    "\n",
    "\n",
    "# Put the first dataset in\n",
    "\n",
    "\n",
    "\n",
    "# List the data \n",
    "# initial Values\n",
    "K1 = 50\n",
    "D1 = 50\n",
    "\n",
    "# for each day, calculate data and insert into db\n",
    "for d in shdaily.find()[:10]:\n",
    "    date = d['date']\n",
    "    datalist = pd.DataFrame(list(shdaily.find({'date':{\"$lte\": date}}).sort('date', -1)))\n",
    "    data = datalist[:9]\n",
    "    \n",
    "    # get previous KDJ data from database\n",
    "    K1 = data.ix[1]['KDJ_K']\n",
    "    D1 = data.ix[1]['KDJ_D']\n",
    "    \n",
    "    high = data['high'].values\n",
    "    low  = data['low'].values\n",
    "    close = data[:1]['close'].values\n",
    "    K1,D1,K2,D2,J2 = KDJCalculation(K1,D1,max(high),min(low),close)\n",
    "    d['KDJ_K'] = K2[0]\n",
    "    d['KDJ_D'] = D2[0]\n",
    "    d['KDJ_J'] = J2[0]\n",
    "    \n",
    "#     K1 = K2\n",
    "#     D1 = D2\n",
    "    print d\n",
    "\n",
    "\n",
    "#datalist = pd.DataFrame(list(shdaily.find().sort('date', -1)))\n",
    "\n",
    "#date1 = datetime.strptime(\"01/01/16\", \"%d/%m/%y\")\n",
    "\n",
    "\n",
    "# List out the data before or equal a specific date\n",
    "#list(shdaily.find({'date':{\"$lte\":'2016-02-08'}}).sort('date', -1))\n",
    "\n",
    "\n",
    "# Get last day KDJ data from database\n",
    "datalist = pd.DataFrame(list(shdaily.find({'date':{\"$lte\": '2016-02-10'}}).sort('date', -1)))\n",
    "data = datalist.ix[1]\n",
    "data['KDJ_K']\n",
    "\n",
    "\n",
    "# Save data to db \n",
    "\n",
    "\n",
    "# data = datalist[:9]\n",
    "\n",
    "# data\n",
    "\n",
    "\n",
    "\n",
    "# K1 = 50\n",
    "# D1 = 50\n",
    "# high = data['high'].values\n",
    "# low  = data['low'].values\n",
    "# close = data[:1]['close'].values\n",
    "\n",
    "# K1,D1,K2,D2,J2 = KDJCalculation(K1,D1,max(high),min(low),close)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Another KDJ Calculation based on dataframe\n",
    "def CalculateKDJ(stock_data):\n",
    "    # Initiate KDJ parameters\n",
    "    endday = pd.datetime.today()\n",
    "    N1= 9\n",
    "    N2= 3\n",
    "    N3= 3\n",
    "    # Perform calculation\n",
    "    #stock_data = get_price(stock, end_date=endday)\n",
    "    low_list = pd.rolling_min(stock_data['LowPx'], N1)\n",
    "    low_list.fillna(value=pd.expanding_min(stock_data['LowPx']), inplace=True)\n",
    "    high_list = pd.rolling_max(stock_data['HighPx'], N1)\n",
    "    high_list.fillna(value=pd.expanding_max(stock_data['HighPx']), inplace=True)\n",
    "    #rsv = (stock_data['ClosingPx'] - low_list) / (high_list - low_list) * 100\n",
    "    \n",
    "    rsv = (stock_data['ClosingPx'] - stock_data['LowPx']) / (stock_data['HighPx'] - stock_data['LowPx']) * 100\n",
    "    stock_data['KDJ_K'] = pd.ewma(rsv, com = N2)\n",
    "    stock_data['KDJ_D'] = pd.ewma(stock_data['KDJ_K'], com = N3)\n",
    "    stock_data['KDJ_J'] = 3 * stock_data['KDJ_K'] - 2 * stock_data['KDJ_D']\n",
    "    KDJ = stock_data[['KDJ_K','KDJ_D','KDJ_J']]\n",
    "    return KDJ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
